import paddle.fluid as fluid

from helm.static.models.layers import Conv2d, Pool2d, GlobalAvgPool, BN, Linear, Act, Layer, Sequential

__all__ = [
    "PyramidNet"
]


class PadChannel(Layer):
    def __init__(self, channels):
        super().__init__()
        assert channels >= 0
        self.c = channels

    def forward(self, x):
        return fluid.layers.pad(x, [0, 0, 0, self.c, 0, 0, 0, 0])

    def extra_repr(self):
        return "padding=%d" % self.c


class Shortcut(Layer):
    def __init__(self, in_channels, out_channels, stride=1):
        super().__init__()
        layers = []
        if stride == 2:
            layers.append(("pool", Pool2d(2, 2)))
        if in_channels != out_channels:
            layers.append(("pad", PadChannel(out_channels - in_channels)))
        self.shortcut = Sequential(*layers)

    def forward(self, x):
        return self.shortcut(x)


class BasicBlock(Layer):
    expansion = 1

    def __init__(self, in_channels, out_channels, stride=1):
        super().__init__()
        self.branch1 = Sequential(
            ("bn1", BN(in_channels)),
            ("conv1", Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, bias=False)),
            ("bn2", BN(out_channels)),
            ("act2", Act()),
            ("conv2", Conv2d(out_channels, out_channels, kernel_size=3, bias=False)),
            ("bn3", BN(out_channels)),
        )
        self.branch2 = Shortcut(in_channels, out_channels, stride)

    def forward(self, x):
        return self.branch1(x) + self.branch2(x)


class Bottleneck(Layer):
    expansion = 4

    def __init__(self, in_channels, channels, stride=1):
        super().__init__()
        out_channels = channels * self.expansion
        branch1 = [
            ("bn1", BN(in_channels)),
            ("conv1", Conv2d(in_channels, channels, kernel_size=1, bias=False)),
            ("bn2", BN(channels)),
            ("act2", Act()),
            ("conv2", Conv2d(channels, channels, kernel_size=3, stride=stride, bias=False)),
            ("bn3", BN(channels)),
            ("act3", Act()),
            ("conv3", Conv2d(channels, out_channels, kernel_size=1, bias=False)),
            ("bn4", BN(out_channels)),
        ]
        self.branch1 = Sequential(*branch1)
        self.branch2 = Shortcut(in_channels, out_channels, stride)

    def forward(self, x):
        return self.branch1(x) + self.branch2(x)


def rd(c):
    return int(round(c, 2))


class PyramidNet(Layer):
    def __init__(self,
                 start_channels,
                 num_classes,
                 block,
                 widening_fractor,
                 depth):
        super().__init__()

        if block == 'basic':
            block = BasicBlock
            num_layers = [(depth - 2) // 6] * 3
        elif block == 'bottleneck':
            block = Bottleneck
            num_layers = [(depth - 2) // 9] * 3
        else:
            raise ValueError("invalid block type: %s" % block)

        strides = [1, 2, 2]

        self.add_channel = widening_fractor / sum(num_layers)
        self.in_channels = start_channels
        self.channels = start_channels

        layers = [("init_block", Conv2d(3, start_channels, kernel_size=3, bn=True))]

        for i, (n, s) in enumerate(zip(num_layers, strides)):
            layers.append(
                ("stage%d" % (i + 1), self._make_layer(block, n, stride=s))
            )

        layers.append(("post_activ", Sequential(
            ("bn", BN(self.in_channels)),
            ("act", Act()),
        )))
        self.features = Sequential(*layers)
        assert (start_channels + widening_fractor) * block.expansion == self.in_channels
        self.final_pool = GlobalAvgPool()
        self.output = Linear(self.in_channels, num_classes)

    def _make_layer(self, block, num_layers, stride):
        self.channels = self.channels + self.add_channel
        layers = [("unit1", block(self.in_channels, rd(self.channels), stride=stride))]
        self.in_channels = rd(self.channels) * block.expansion
        for i in range(1, num_layers):
            self.channels = self.channels + self.add_channel
            layers.append(("unit%d" % (i + 1), block(self.in_channels, rd(self.channels))))
            self.in_channels = rd(self.channels) * block.expansion
        return Sequential(*layers)

    def forward(self, x):
        x = self.features(x)
        x = self.final_pool(x)
        x = self.output(x)
        return x
